import numpy as np
import torch
import random
import torch.nn as nn
import matplotlib.pyplot as plt
import torch.nn.init as init

# 1. 创造数据，数据集
point = [[0.5, 8.6], [0.5, 9.3], [0.6, 8.9], [0.6, 8.3], [0.6, 8.0], [0.7, 7.8], [0.7, 8.9], [0.7, 9.7],
              [0.7, 9.1], [0.8, 9.2],
              [0.8, 8.5], [0.8, 8.4], [0.9, 8.8], [0.9, 8.6], [0.9, 8.2], [1.0, 8.2], [1.0, 6.6], [1.0, 6.3],
              [1.0, 6.9], [1.1, 7.1],
              [1.1, 7.7], [1.1, 6.5], [1.2, 7.0], [1.2, 7.7], [1.2, 6.1], [1.3, 7.7], [1.3, 6.5], [1.3, 6.9],
              [1.3, 5.3], [1.4, 5.7],
              [1.4, 5.8], [1.4, 5.6], [1.5, 6.8], [1.5, 6.7], [1.5, 6.6], [1.6, 3.6], [1.6, 5.3], [1.6, 6.9],
              [1.6, 5.9], [1.7, 6.0],
              [1.7, 4.7], [1.7, 5.0], [1.8, 4.5], [1.8, 5.6], [1.8, 4.2], [1.9, 3.8], [1.9, 4.5], [1.9, 5.8],
              [1.9, 6.7], [2.0, 6.5],
              [2.0, 6.3], [2.0, 4.9], [2.1, 5.9], [2.1, 3.6], [2.1, 3.8], [2.2, 4.8], [2.2, 4.3], [2.2, 4.6],
              [2.2, 4.1], [2.3, 3.5],
              [2.3, 2.9], [2.3, 4.4], [2.4, 4.5], [2.4, 3.6], [2.4, 4.3], [2.5, 5.0], [2.5, 2.3], [2.5, 4.4],
              [2.5, 6.0], [2.6, 3.4],
              [2.6, 3.6], [2.6, 3.6], [2.7, 4.9], [2.7, 3.6], [2.7, 5.1], [2.8, 5.1], [2.8, 3.5], [2.8, 2.0],
              [2.8, 3.7], [2.9, 2.5],
              [2.9, 3.3], [2.9, 2.8], [3.0, 2.5], [3.0, 1.4], [3.0, 4.1], [3.1, 2.8], [3.1, 4.1], [3.1, 2.2],
              [3.1, 3.1], [3.2, 3.2],
              [3.2, 3.0], [3.2, 3.7], [3.3, 3.7], [3.3, 2.9], [3.3, 4.0], [3.4, 2.7], [3.4, 3.0], [3.4, 2.3],
              [3.4, 1.8], [3.5, 3.4],
              [3.5, 3.9], [3.5, 3.1], [3.6, 3.1], [3.6, 2.4], [3.6, 2.1], [3.7, 2.3], [3.7, 1.3], [3.7, 2.7],
              [3.8, 2.0], [3.8, 2.2],
              [3.8, 3.0], [3.8, 2.0], [3.9, 3.1], [3.9, 1.9], [3.9, 0.0], [4.0, 1.6], [4.0, 1.9], [4.0, 1.8],
              [4.1, 2.6], [4.1, 2.0],
              [4.1, 1.2], [4.1, 2.5], [4.2, 2.0], [4.2, 0.1], [4.2, 1.7], [4.3, 1.2], [4.3, 2.4], [4.3, 2.1],
              [4.4, 1.3], [4.4, 1.0],
              [4.4, 1.6], [4.4, 2.8], [4.5, 2.8], [4.5, 2.1], [4.5, 1.9], [4.6, 3.0], [4.6, 2.3], [4.6, 2.3],
              [4.7, 3.0], [4.7, 0.4],
              [4.7, 1.6], [4.7, 1.1], [4.8, 2.6], [4.8, 2.9], [4.8, 2.9], [4.9, 2.5], [4.9, 2.4], [4.9, 1.9],
              [5.0, 1.9], [5.0, 2.9],
              [5.0, 1.4], [5.0, 2.0], [5.1, 3.4], [5.1, 2.5], [5.1, 1.7], [5.2, 2.7], [5.2, 2.2], [5.2, 1.9],
              [5.3, 1.5], [5.3, 2.6],
              [5.3, 1.9], [5.3, 1.2], [5.4, 2.2], [5.4, 2.6], [5.4, 1.2], [5.5, 1.8], [5.5, 2.4], [5.5, 3.0],
              [5.6, 2.7], [5.6, 3.6],
              [5.6, 2.2], [5.6, 2.4], [5.7, 2.2], [5.7, 3.3], [5.7, 2.2], [5.8, 3.0], [5.8, 0.9], [5.8, 2.6],
              [5.9, 2.5], [5.9, 1.5],
              [5.9, 2.4], [5.9, 2.1], [6.0, 2.2], [6.0, 1.7], [6.0, 2.8], [6.1, 1.4], [6.1, 2.5], [6.1, 2.0],
              [6.2, 2.5], [6.2, 2.5],
              [6.2, 1.0], [6.2, 2.4], [6.3, 1.1], [6.3, 2.9], [6.3, 3.5], [6.4, 2.3], [6.4, 5.0], [6.4, 2.8],
              [6.5, 1.5], [6.5, 4.0],
              [6.5, 3.6], [6.6, 3.8], [6.6, 2.7], [6.6, 2.6], [6.6, 2.1], [6.7, 3.1], [6.7, 3.6], [6.7, 3.5],
              [6.8, 2.7], [6.8, 3.0],
              [6.8, 2.5], [6.9, 2.9], [6.9, 3.9], [6.9, 3.6], [6.9, 3.4], [7.0, 3.4], [7.0, 3.4], [7.0, 4.5],
              [7.1, 3.9], [7.1, 4.6],
              [7.1, 4.4], [7.2, 4.1], [7.2, 3.2], [7.2, 2.7], [7.2, 4.2], [7.3, 4.1], [7.3, 5.7], [7.3, 3.7],
              [7.4, 3.0], [7.4, 4.0],
              [7.4, 3.9], [7.5, 4.3], [7.5, 3.5], [7.5, 4.4], [7.5, 6.2], [7.6, 4.0], [7.6, 5.7], [7.6, 6.6],
              [7.7, 6.1], [7.7, 5.4],
              [7.7, 2.5], [7.8, 5.6], [7.8, 4.1], [7.8, 5.9], [7.8, 5.1], [7.9, 4.5], [7.9, 5.1], [7.9, 5.5],
              [8.0, 5.8], [8.0, 5.0],
              [8.0, 6.0], [8.1, 5.8], [8.1, 5.9], [8.1, 5.6], [8.1, 5.2], [8.2, 4.0], [8.2, 6.4], [8.2, 4.5],
              [8.3, 6.2], [8.3, 5.7],
              [8.3, 5.3], [8.4, 4.9], [8.4, 6.9], [8.4, 5.0], [8.4, 7.4], [8.5, 5.0], [8.5, 7.5], [8.5, 7.1],
              [8.6, 6.4], [8.6, 6.0],
              [8.6, 7.5], [8.7, 5.8], [8.7, 7.7], [8.7, 6.2], [8.7, 6.6], [8.8, 6.2], [8.8, 8.1], [8.8, 7.7],
              [8.9, 7.4], [8.9, 8.2],
              [8.9, 7.4], [9.0, 7.6], [9.0, 6.7], [9.0, 7.7], [9.0, 8.2], [9.1, 7.7], [9.1, 9.2], [9.1, 9.1],
              [9.2, 8.5], [9.2, 7.4],
              [9.2, 8.5], [9.3, 9.2], [9.3, 8.3], [9.3, 9.7], [9.3, 8.5], [9.4, 8.2], [9.4, 9.9], [9.4, 8.5],
              [9.5, 9.9], [9.5, 8.7]]

# 将原始数据分成训练和测试两类
random.shuffle(point)
# 取前15%作为测试数据集
split_index = int(0.15 * len(point))
train_point = np.array(point[split_index:])
test_point = np.array(point[:split_index])
# 训练数据集当中的训练数据集
x_train = train_point[:, 0]
# 训练数据集当中的验证数据集
y_train = train_point[:, 1]
# 验证数据集当中的训练数据
x_test = test_point[:, 0]
# 验证数据集当中的验证数据
y_test = test_point[:, 1]

"""
tensorflow里面用的是张量 => 所以我们要转换成张量
"""
# 将Numpy数组转成tensor
x_train_tensor = torch.tensor(x_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
x_test_tensor = torch.tensor(x_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32)

# 2.定义前向模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 虚拟仿真里面用的就是用
        self.layer1 = nn.Linear(1, 6)
        self.layer2 = nn.Linear(6, 1)

        # 初始化方式
        # 用init这个包专门来初始化参数
        init.xavier_normal_(self.layer1.weight) # 初始化权重
        init.zeros_(self.layer1.bias) # 初始化偏置

    def forward(self, x):
        x = torch.tanh(self.layer1(x))
        x = self.layer2(x)
        return x

model = Model()

# 3.定义损失函数和优化器
cri = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# 初始化绘图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
train_loss_list, test_loss_list, train_epoch_list, test_epoch_list = [], [], [], []

# 4.开始迭代
epoches = 5000
for epoch in range(1, epoches+1):
    """
    前向传播
        调用模型 => 创建一个线性层，它期望每个输入样本有 1 个特征，然后输出 6 个特征 nn.Linear(1, 6) 里的这个 1，就是问题的根源！
        原始数据 x_train_tensor 形状是 [170]  [ 0.5, 0.6, 0.7, 0.8, ... ]  170个数字
        但 模型期望的是 =>
              [[0.5],  # 样本1：特征值是0.5
               [0.6],  # 样本2：特征值是0.6
               [0.7],  # 样本3：特征值是0.7
               ...]
            所以 => .unsqueeze(1)
    """
    y_pred = model(x_train_tensor.unsqueeze(1))
    loss = cri(y_pred.squeeze(1), y_train_tensor)  # 训练的损失

    # 更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 5.显示频率设置
    if epoch == 1 or epoch % 100 == 0:
        # 6.绘图
        print(f"epoch: {epoch}, loss: {loss}")
        # 更新完才拿到test损失
        y_pred_test = model(x_test_tensor.unsqueeze(1))
        test_loss = cri(y_pred_test.squeeze(1), y_test_tensor)
        train_loss_list.append(loss.detach().numpy())
        test_loss_list.append(test_loss.detach().numpy())
        train_epoch_list.append(epoch)
        test_epoch_list.append(epoch)

        x_range = torch.tensor(np.linspace(0, 10, 100), dtype=torch.float32).unsqueeze(1)
        y_range = model(x_range).detach().numpy()

        # 绘制左侧拟合图
        ax1.clear()
        ax1.scatter(x_train, y_train)
        ax1.plot(x_range, y_range, c="r")

        ax2.clear()
        ax2.plot(train_epoch_list, train_loss_list, "r", label="train loss")
        ax2.plot(test_epoch_list, test_loss_list, "b", label="test loss")
        ax2.legend()

        plt.pause(0.2)
plt.show()